from service.rag_neo4j_service import RagNeo4jService
import json


class PersonNodeLLMService:

    def __init__(self):
        self.ragNeo4jService = RagNeo4jService()
        self.chain = self.ragNeo4jService.create_chain()

    # 处理json格式的字符串
    def handle_json(self, json_str):
        if "```json" in json_str:
            text = json_str.split("```json")[1].split("```")[0]
        else:
            text = json_str
        # 转换成json格式
        info = json.loads(text)
        return info

    # 创建节点
    def create_node(self, **kwargs):
        info = f"创建一个Person节点,它的属性有：{kwargs}"
        self.chain.invoke({"query": info})

    # 查询节点
    def query_node(self, **kwargs):
        if len(kwargs) > 0:
            info = f"""
                    查询符合条件的人员节点（高管和个人股东），生成正确的 Neo4j Cypher 语句，需满足：
                    1. 子查询1：查询标签为"高管"的节点，返回字段：name（姓名）、position（职位）、tenure（任职时间）、type（固定为"高管"）
                    2. 子查询2：查询标签为"十大股东"且股本性质为"个人"的节点，返回字段：name（股东名称）、position（固定为"个人股东"）、share_ratio（持股比例）、type（固定为"个人股东"）
                    3. 用 UNION ALL 合并两个子查询结果（确保两个子查询返回字段数量和类型一致）
                    4. 合并后用 WITH 收集所有结果为 person_list，再调用 apoc.convert.toTree(person_list, false) 转换为 JSON 格式
                    5. 最终 RETURN value（APOC 转换后的结果）
                    6. 严格遵守 Neo4j 语法：UNION 只能连接两个完整子查询，WITH/CALL 必须在 UNION 之后
                    符合条件的属性有：{kwargs}
                    """
        else:
            info = f"""
                    查询所有人员节点（包括高管和个人股东），生成正确的 Neo4j Cypher 语句，需满足：
                    1. 子查询1：MATCH (p:高管) RETURN p.name AS name, p.职位 AS position, p.任职时间 AS tenure, '高管' AS type
                    2. 子查询2：MATCH (s:十大股东 {{股本性质: '个人'}}) RETURN s.股东名称 AS name, '个人股东' AS position, s.持股比例 AS share_ratio, '个人股东' AS type
                    3. 用 UNION ALL 合并两个子查询结果
                    4. 合并后执行：WITH collect({{name:name, position:position, extra:CASE WHEN type='高管' THEN tenure ELSE share_ratio END, type:type}}) AS person_list
                    5. 再执行：CALL apoc.convert.toTree(person_list, false) YIELD value
                    6. 最后 RETURN value
                    7. 禁止在 UNION 子查询中间插入 RETURN 以外的语句，严格遵守 Neo4j 语法
                    """
        try:
            rs = self.chain.invoke({"query": info})
            # 打印大模型生成的Cypher，方便排查
            print("大模型生成的Cypher：", rs.get("generated_cypher", "无"))
            data = self.handle_json(rs["result"])
            print("查询结果：", data)
            return data
        except Exception as e:
            print(f"查询节点失败: {str(e)}")
            raise  # 抛出异常让视图层捕获

    # 删除节点
    def delete_node(self, **kwargs):
        info = f"删除一个Person节点,它的属性有：{kwargs}"
        self.chain.invoke({"query": info})

    # 更新节点
    def update_node(self, **kwargs):
        info = f"更新一个Person节点,它的属性有：{kwargs}"
        self.chain.invoke({"query": info})


if __name__ == '__main__':
    personNodeLLMService = PersonNodeLLMService()
    # personNodeLLMService.create_node(name="王六",age=18,city="成都",occupation="学生")
    # personNodeLLMService.query_node(name="王六")
    # personNodeLLMService.delete_node(name="王六")
    # personNodeLLMService.update_node(name="李四", age=19)
